import sys
import numpy as np
import pyaudio

sys.path.insert(0, r'C:\tools\projects\python\GUI\test1\config\third_party\Matcha-TTS')
print(sys.path)
from cosyvoice.cli.cosyvoice import CosyVoice2
from cosyvoice.utils.file_utils import load_wav
import torchaudio

cosyvoice = CosyVoice2('pretrained_models/CosyVoice2-0.5B', load_jit=False, load_trt=False)
print(cosyvoice.sample_rate)
prompt_speech_16k = load_wav('D:\\tools\python\CosyVoice\q3.wav', 16000)
prompt_text = '哈哈哈，笑死我啦，假仁门谁懂啊，今天遇到了一个虾头男'
p = pyaudio.PyAudio()
stream = p.open(
    format=pyaudio.paInt16,
    channels=1,
    rate=cosyvoice.sample_rate,
    output=True
)
voice = None
for i, j in enumerate(cosyvoice.inference_zero_shot(
        '收到好友从远方寄来的生日礼物，那份意外的惊喜与深深的祝福让我心中充满了甜蜜的快乐，笑容如花儿般绽放。',
        prompt_text, prompt_speech_16k, stream=False)):
    # 实时播放音频块
    audio_tensor = j['tts_speech']
    audio_array = audio_tensor.numpy()
    if audio_array.ndim > 1:
        audio_array = audio_array[0]
    audio_array = (audio_array * 32767).astype(np.int16)
    stream.write(audio_array.tobytes())
    torchaudio.save('zero_shot_{}.wav'.format(i), j['tts_speech'], cosyvoice.sample_rate)
# 清理资源
stream.stop_stream()
stream.close()
p.terminate()
